package com.zdb.demo.flink.streaming;

import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.streaming.api.TimeCharacteristic;
import org.apache.flink.streaming.api.datastream.AsyncDataStream;
import org.apache.flink.streaming.api.datastream.DataStream;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.async.ResultFuture;
import org.apache.flink.streaming.api.functions.async.RichAsyncFunction;

import java.io.BufferedReader;
import java.io.FileReader;
import java.util.Collections;
import java.util.concurrent.Callable;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.Future;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;

public class AsyncDemo {

    public static void main(String[] args) throws Exception {
        Configuration conf = new Configuration(){{
            setInteger("rest.port", 9191);
            setBoolean("local.start-webserver", true);
        }};
        final StreamExecutionEnvironment streamEnv = StreamExecutionEnvironment.createLocalEnvironmentWithWebUI(conf);
        streamEnv.setParallelism(4);

        streamEnv.setStreamTimeCharacteristic(TimeCharacteristic.IngestionTime);
        DataStream<String> userNameDs = streamEnv.fromElements("a", "b", "c", "d", "a", "b", "c", "d", "a", "b", "c", "d", "a", "b", "c", "d");

        AsyncDataStream.orderedWait(userNameDs,
                new AsyncIORequest(), 5000, TimeUnit.MILLISECONDS, 100).print();

        streamEnv.execute("stream async");
    }
}

class AsyncIORequest extends RichAsyncFunction<String, Tuple2<String, String>> {

    ExecutorService executor = null;

    @Override
    public void open(Configuration parameters) throws Exception {
        executor = Executors.newCachedThreadPool();
    }

    @Override
    public void close() throws Exception {
        executor.shutdown();
    }

    @Override
    public void asyncInvoke(String input, ResultFuture<Tuple2<String, String>> resultFuture) throws Exception {
        Future<String> result = executor.submit(new Callable<String>() {
            public String call() throws Exception {
                BufferedReader br = new BufferedReader(new FileReader("E:\\logs\\tmp\\user_profile.csv"));
                String profileName = "";

                String line = null;
                while ((line = br.readLine()) != null) {
                    String[] kv = line.split(",");
                    if(input.equals(kv[0])) {
                        profileName = kv[1];
                        break;
                    }
                }
                br.close();
                return profileName;
            }
        });

        CompletableFuture.supplyAsync(() -> {
            try {
                return  result.get();
            } catch (InterruptedException | ExecutionException e) {
                e.printStackTrace();
                return null;
            }
        }).thenAccept((String profileName) -> {
            resultFuture.complete(Collections.singleton(new Tuple2<>(input, profileName)));
        });
    }

    @Override
    public void timeout(String input, ResultFuture<Tuple2<String, String>> resultFuture) throws Exception {
        resultFuture.completeExceptionally(new TimeoutException("Async fuction call has ts out."));
    }
}
